import requests
import os
from urllib.parse import urlparse
from tqdm import tqdm

class MP3Downloader:
    """MP3文件下载器类，支持指定保存路径和显示下载进度"""
    
    def __init__(self, save_dir=None):
        """
        初始化下载器
        
        参数:
            save_dir: 保存MP3文件的目录，默认为当前工作目录
        """
        # 设置默认保存目录
        self.save_dir = save_dir if save_dir else os.getcwd()
        # 确保保存目录存在
        self._ensure_dir_exists()
        
        # 设置请求头，模拟浏览器访问
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
    
    def _ensure_dir_exists(self):
        """确保保存目录存在，不存在则创建"""
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir, exist_ok=True)
    
    def set_save_dir(self, new_dir):
        """
        修改保存目录
        
        参数:
            new_dir: 新的保存目录路径
        """
        self.save_dir = new_dir
        self._ensure_dir_exists()
    
    def _get_filename_from_url(self, url):
        """
        从URL中提取文件名
        
        参数:
            url: MP3文件的URL
            
        返回:
            提取的文件名，如果无法提取则返回默认名称
        """
        parsed_url = urlparse(url)
        filename = os.path.basename(parsed_url.path)
        
        # 如果文件名不包含.mp3扩展名，添加默认扩展名
        if not filename.endswith('.mp3'):
            filename = f"downloaded_audio_{hash(url)}.mp3"
            
        return filename
    
    def download(self, url, filename=None, show_progress=True):
        """
        下载MP3文件
        
        参数:
            url: MP3文件的URL
            filename: 可选，指定保存的文件名，不指定则从URL提取
            show_progress: 是否显示下载进度条
            
        返回:
            成功则返回保存的文件路径，失败则返回None
        """
        try:
            # 获取文件名
            if not filename:
                filename = self._get_filename_from_url(url)
            
            # 构建完整保存路径
            save_path = os.path.join(self.save_dir, filename)
            
            # 检查文件是否已存在
            if os.path.exists(save_path):
                print(f"文件已存在: {save_path}")
                return save_path
            
            # 发送请求
            with requests.get(url, headers=self.headers, stream=True, timeout=10) as response:
                response.raise_for_status()  # 检查请求是否成功
                
                # 获取文件大小
                file_size = int(response.headers.get('content-length', 0))
                
                # 写入文件
                with open(save_path, 'wb') as f:
                    if show_progress and file_size > 0:
                        # 带进度条的写入
                        with tqdm(total=file_size, unit='B', unit_scale=True, 
                                 desc=filename, ncols=100) as pbar:
                            for chunk in response.iter_content(chunk_size=8192):
                                if chunk:  # 过滤掉保持连接的空块
                                    f.write(chunk)
                                    pbar.update(len(chunk))
                    else:
                        # 不带进度条的写入
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                        print(f"下载完成: {filename}")
            
            return save_path
            
        except requests.exceptions.RequestException as e:
            print(f"下载失败: {str(e)}")
            # 清理可能的不完整文件
            if os.path.exists(save_path):
                os.remove(save_path)
            return None
        except Exception as e:
            print(f"发生错误: {str(e)}")
            return None

# 使用示例
if __name__ == "__main__":
    # 创建下载器实例，指定保存目录
    downloader = MP3Downloader(save_dir="./downloaded_mp3")
    
    # 示例MP3 URL（实际使用时替换为有效的MP3链接）
    mp3_url = "https://er-sycdn.kuwo.cn/9b970fbc2dbd79f72d5311d18068fc9a/68dbd485/resource/30106/trackmedia/C200002PbOJa1LHVnv.m4a?from=vip"  # 请替换为实际的MP3 URL
    
    # 下载文件
    downloaded_path = downloader.download(mp3_url)
    
    if downloaded_path:
        print(f"文件已保存至: {downloaded_path}")
    
    # 可以修改保存目录
    # downloader.set_save_dir("./new_mp3_folder")
    # 再次下载
    # downloader.download("https://example.com/another_song.mp3", filename="my_song.mp3")